#!/usr/bin/env python3

# strict_pam - manipulating PAM configuration in a strict manner
#
# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import abc
import enum
import pathlib
import functools
from ._util import Util
from ._error import ParseError, VerifyError, FeatureNotFoundError


PAM_SERVICE_NAME_OTHER = "other"


class PamGroup(enum.Enum):
    AUTH = "auth"
    ACCOUNT = "account"
    PASSWORD = "password"
    SESSION = "session"

    def __str__(self):
        return self.value


class PamControl(enum.Enum):
    REQUISITE = "requisite"
    REQUIRED = "required"
    SUFFICIENT = "sufficient"
    OPTIONAL = "optional"
    INCLUDE = "include"
    SUBSTACK = "substack"

    def __str__(self):
        return self.value


STANDARD_PAM_GROUPS = [PamGroup.AUTH, PamGroup.ACCOUNT, PamGroup.PASSWORD, PamGroup.SESSION]


STANDARD_PAM_CONTROLS = [PamControl.REQUISITE, PamControl.REQUIRED, PamControl.SUFFICIENT, PamControl.OPTIONAL]


class Context:

    def __init__(self, dirPrefix="/"):
        self._dirPrefix = dirPrefix
        self._pamDir = os.path.join(self._dirPrefix, "etc", "pam.d")
        self._serviceCache = dict()

    @property
    def dir_prefix(self):
        return self._dirPrefix

    @functools.cache
    def get_service_names(self):
        ret = []

        # get data
        for fn in os.listdir(self._pamDir):
            if fn.startswith("system-"):                   # filter library services
                continue
            ret.append(fn)

        # sort data, service "other" should be the last
        ret.sort()
        try:
            ret.remove(PAM_SERVICE_NAME_OTHER)
            ret.append(PAM_SERVICE_NAME_OTHER)
        except ValueError:
            pass

        return ret

    def get_service(self, service_name):
        if service_name not in self._serviceCache:
            self._serviceCache[service_name] = PamService(self, service_name)
        return self._serviceCache[service_name]


class PamService:

    def __init__(self, context, service_name):
        self._context = context
        self._serviceName = service_name

        # file path
        self._fullfn = os.path.join(self._context._pamDir, self._serviceName)
        if not os.path.isfile(self._fullfn):
            raise ParseError("file does not exist", self._fullfn)

        # content
        self._content = pathlib.Path(self._fullfn).read_text()
        self._data = dict()           # dict<pam-group, pam-rule-list>
        self._bComplex = False        # we're lucky that we haven't encountered complex pam service file yet
        self._basedOn = None
        if True:
            for i, line in enumerate(self._content.split("\n")):
                line = re.sub("#.*", "", line.strip())
                if line == "":
                    continue

                m = re.fullmatch(r"(\S+)\s+(\S+)\s+(\S+)(\s+.*)?", line)
                if m is None:
                    raise ParseError("illegal format", self._fullfn, i + 1)
                fType = m.group(1)
                fControl = m.group(2)
                fModulePath = m.group(3)
                fModuleArguments = m.group(4)

                try:
                    fType = PamGroup(fType)
                except ValueError:
                    raise ParseError("\"%s\" is invalid" % (fType), self._fullfn, i + 1)

                try:
                    fControl = PamControl(fControl)
                except ValueError:
                    raise ParseError("\"%s\" is invalid" % (fControl), self._fullfn, i + 1)

                if fModuleArguments is not None:
                    fModuleArguments = fModuleArguments.lstrip()

                if fControl == PamControl.SUBSTACK:
                    raise ParseError("\"%s\" is not recommended" % (fControl), self._fullfn, i + 1)

                if fControl == PamControl.INCLUDE:
                    if fModuleArguments is not None:
                        raise ParseError("module arguments should not exist for \"%s\"" % (fControl), self._fullfn, i + 1)
                    if self._basedOn is not None:
                        if self._basedOn != fModulePath:
                            raise ParseError("multiple PAM services (\"%s\", \"%s\") get included" % (self._basedOn, fModulePath), self._fullfn, i + 1)
                    else:
                        self._basedOn = fModulePath

                Util.dictPut(self._data, fType, []).append(PamRule(fControl, fModulePath, fModuleArguments))

        # verify based-on
        if self._basedOn is not None:
            for fType, ruleList in self._data.items():
                if len([x for x in ruleList if x.control == PamControl.INCLUDE]) != 1:
                    raise VerifyError("there must be one and only one \"include\" in PAM group \"%s\"" % (fType), self._fullfn)
                if fType != PamGroup.SESSION and ruleList[-1].control != PamControl.INCLUDE:
                    raise VerifyError("\"include\" is not the last rule in PAM group \"%s\"" % (fType), self._fullfn)

        # verify control auth and account
        if PamGroup.AUTH in self._data.keys():
            if PamGroup.ACCOUNT not in self._data.keys():
                raise VerifyError("PAM group %s and %s must exist simultaneously" % (PamGroup.AUTH, PamGroup.ACCOUNT), self._fullfn)
        else:
            if PamGroup.ACCOUNT in self._data.keys():
                raise VerifyError("PAM group %s and %s must exist simultaneously" % (PamGroup.AUTH, PamGroup.ACCOUNT), self._fullfn)

        # verify sufficient is in front
        for ruleList in self._data.values():
            bEnd = False
            for rule in ruleList:
                if not bEnd:
                    if rule.control != PamControl.SUFFICIENT:
                        bEnd = True
                else:
                    if rule.control == PamControl.SUFFICIENT:
                        raise VerifyError("control \"%s\" should always be in front of other controls" % (PamControl.SUFFICIENT), self._fullfn)

        # verify should not use
        if True:
            modulePathList = [
                "pam_env.so",            # pam_env.so is complex and we have the ability not to use it
                "pam_mail.so",           # pam_mail.so is really ancient
                "pam_limits.so",         # pam_limits.so is complex and we have the ability not to use it
                "pam_cap.so",            # pam_cap.so is complex and we do not know what it is used for
                "pam_group.so",          # pam_group.so seems more complex than we imagine
                "pam_namespace.so",      # we don't know pam_namespace.so
            ]
            for ruleList in self._data.values():
                for rule in ruleList:
                    if rule.module_path in modulePathList:
                        raise VerifyError("%s should not be used" % (rule.module_path), self._fullfn)

        # verify common empty arguments
        if True:
            modulePathList = [
                "pam_permit.so",
                "pam_deny.so",
            ]
            for ruleList in self._data.values():
                for rule in ruleList:
                    if rule.module_path in modulePathList and rule.module_arguments is not None:
                        raise VerifyError("%s should not have any argument" % (rule.module_path), self._fullfn)

        # verify pam service "other"
        if service_name == PAM_SERVICE_NAME_OTHER:
            if set(self._data.keys()) != set(STANDARD_PAM_GROUPS):
                raise VerifyError("invalid content", self._fullfn)
            for ruleList in self._data.values():
                if len(ruleList) != 1:
                    raise VerifyError("invalid content", self._fullfn)
                if ruleList[0].control != PamControl.REQUIRED:
                    raise VerifyError("invalid content", self._fullfn)
                if ruleList[0].module_path != "pam_deny.so":
                    raise VerifyError("invalid content", self._fullfn)

        # verify modules
        for fType, ruleList in self._data.items():
            for rule in ruleList:
                if rule.control == PamControl.INCLUDE:
                    continue
                if not os.path.exists("/lib64/security/" + rule.module_path):
                    raise VerifyError("non-exist module \"%s\"" % (rule.module_path), self._fullfn)
                if True:
                    ret = Util.pamGetModuleTypesProvided(rule.module_path)
                    if ret is not None and str(fType).replace("-", "") not in ret:
                        raise VerifyError("module \"%s\" is not suitable for %s" % (rule.module_path, fType), self._fullfn)

        # features
        self._features = []
        if not self._bComplex:
            for featureCls in FeatureBase.__subclasses__():
                if not featureCls.__name__.startswith('_'):
                    try:
                        self._features.append(featureCls(self._context, self))
                    except FeatureNotFoundError:
                        pass
            for f in self._features:
                f.verify_after_init()

        # verify features
        pass

    @property
    def name(self):
        return self._serviceName

    @property
    def file_path(self):
        return self._fullfn

    @property
    def content(self):
        return self._format(self._data)

    @property
    def groups(self):
        return Util.oIterate(self._data.keys(), STANDARD_PAM_GROUPS)

    def is_complex(self):
        return self._bComplex

    def get_based_on(self):
        assert not self._bComplex

        return self._basedOn

    def get_features(self, need_cls=False):
        assert not self._bComplex

        # features are un-ordered
        if need_cls:
            return set([x.__class__ for x in self._features])
        else:
            return set(self._features)

    def get_feature(self, feature_cls):
        assert not self._bComplex

        for f in self._features:
            if f.__class__ == feature_cls:
                return f
        return None

    def add_feature(self, feature_cls, *kargs, **kwargs):
        assert self.get_feature(feature_cls) is None

        newData = self._data.copy()
        featureObj = feature_cls(self._context, self, *kargs, *kwargs)
        featureObj.modify(newData, True)
        self._data = newData

        self._features.append(featureObj)

    def remove_feature(self, feature_cls):
        assert not self._bComplex

        feature_obj, feature_index = None, None
        for i, f in enumerate(self._features):
            if f.__class__ == feature_cls:
                feature_obj, feature_index = f, i
                break
        if feature_obj is None:
            return

        newData = self._data.copy()
        feature_obj.modify(newData, False)
        self._data = newData

        self._features.pop(feature_index)

    def remove_all_features(self):
        assert not self._bComplex

        for fObj in reversed(self._features):
            newData = self._data.copy()
            fObj.modify(newData, False)
            self._data = newData
            self._features.pop()

    def save(self, force=False):
        buf = self._format(self._data)
        if force or buf != self._content:
            pathlib.Path(self._fullfn).write_text(buf)

    @staticmethod
    def _format(data):
        ret = ""
        for mg in Util.oIterate(data.keys(), STANDARD_PAM_GROUPS):
            for item in data[mg]:
                if len(item) == 2:
                    ret += "%17s%17s%s\n" % (mg, item.control, item.module_path)
                elif len(item) == 3:
                    ret += "%17s%17s%s %s\n" % (mg, item.control, item.module_path, item.module_arguments)
                else:
                    assert False
        return ret


class PamRule:

    def __init__(self, control, module_path, module_arguments=None):
        self.control = control
        self.module_path = module_path
        self.module_arguments = module_arguments


class FeatureBase(abc.ABC):

    def __init__(self, context, pam_service):
        self._context = context
        self._pamService = pam_service

    @staticmethod
    @abc.abstractmethod
    def name():
        pass

    @abc.abstractmethod
    def verify_after_init(self):
        pass

    @abc.abstractmethod
    def modify(self, target_data, add_or_remove):
        pass
